#!/usr/bin/env python
# -*- coding=utf-8 -*-

#coordinations in avinput 
# Chr     Start   End     Ref     Alt

#coordinations in vcf
# CHROM   POS     REF     ALT

#======avinput to vcf========
#---snp conversion---- 
# Start to POS
# Delete End 

#---indel conversion----
# Ref Alt
# C - ==> CC C # deletion
# - C ==> C  CC # insertion
# POS = Start-1
# Delete End

import argparse
import os

from pysam import FastaFile

from utils import TitleParser,gopen

class avinputVCFConverter():

    def __init__(self,fasta):
        self.fa = FastaFile(fasta)

    def avinput2vcf(self,avsite):
        """
        avsite:(Chr     Start   End     Ref     Alt)
        Chr:str
        Start/End:int
        """
        chrom,start,_,ref,alt = avsite
        chrom = chrom.replace("chr","")
        
        if ref == "-": #insertion
            pos = start
            region = "{chrom}:{start}-{end}".format(
                chrom=chrom,start=pos,end=pos,
            )
            prev_base = self.fa.fetch(region=region)
            ref =  prev_base 
            alt = prev_base + alt

        elif alt == "-": # deletion
            pos = start -1
            region = "{chrom}:{start}-{end}".format(
            chrom=chrom,start=pos,end=pos,
            )
            prev_base = self.fa.fetch(region=region)
            ref = prev_base + ref
            alt = prev_base 
            
        else: #snp
            pos = start

        return (chrom,pos,ref,alt)
    
    def avinput_file_to_tab_file(self,infile,ofile):
        """
        reading in a avinput file and
        replace the (Chr Start End REF ALT) to (Chrom POS Ref Alt)
        """
        with gopen(infile,'r') as indata,\
            gopen(ofile,'w') as odata:
            title = indata.readline()
            tp = TitleParser(title)
            title_list = title.strip().split('\t')
            del title_list[tp.get_idx('end')]
            title_list[title_list.index("Chr")] = "Chrom"
            title_list[title_list.index("Start")] = "POS"
            odata.write("\t".join(title_list) + "\n")
            for line in indata:
                line_list = line.strip().split('\t')
                chrom = tp.get_field(line_list,'chr',check=False)
                start = int(tp.get_field(line_list,'start',check=False))
                end = int(tp.get_field(line_list,'end',check=False))
                ref = tp.get_field(line_list,'ref',check=False)
                alt = tp.get_field(line_list,'alt',check=False)
                chrom,pos,ref,alt = self.avinput2vcf((chrom,start,end,ref,alt))
                line_list[tp.get_idx('chr')] = chrom
                line_list[tp.get_idx('start')] = str(pos)
                line_list[tp.get_idx('ref')] = ref
                line_list[tp.get_idx('alt')] = alt
                del line_list[tp.get_idx('end')]
                odata.write("\t".join(line_list) + "\n")
        print('Converted file output to {ofile}'.format(ofile=ofile))

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--infile','-i',help="avinput file")
    parser.add_argument("--ofile","-o",help="ofile")
    parser.add_argument("--fasta",'-fa',default="/PUBLIC/database/HUMAN/genome/Human/human_g1k_v37_decoy.fasta",help="fasta file")
    args = vars(parser.parse_args())
    
    converter = avinputVCFConverter(args.get("fasta"))
    converter.avinput_file_to_tab_file(args.get("infile"),args.get("ofile"))
